import torch.nn as nn

# 定义卷积神经网络模型
class UNet1d(nn.Module):
    def __init__(self, num_input_channels=3, feature_scale=4, num_output_channels=3):
        super(UNet1d, self).__init__()
        self.feature_scale = feature_scale
        filters = [64, 128, 256, 512, 1024]
        filters = [x // self.feature_scale for x in filters]

        self.start = nn.Conv1d(in_channels=num_input_channels, out_channels=filters[0], kernel_size=3,padding=1)
        self.down1 = nn.Conv1d(in_channels=filters[0], out_channels=filters[1], kernel_size=3,padding=1)
        self.down2 = nn.Conv1d(in_channels=filters[1], out_channels=filters[2], kernel_size=3,padding=1)
        self.down3 = nn.Conv1d(in_channels=filters[2], out_channels=filters[3], kernel_size=3,padding=1)
        self.down4 = nn.Conv1d(in_channels=filters[3], out_channels=filters[4], kernel_size=3,padding=1)

        self.up4 = nn.Conv1d(in_channels=filters[4], out_channels=filters[3], kernel_size=3,padding=1)
        self.up3 = nn.Conv1d(in_channels=filters[3], out_channels=filters[2], kernel_size=3,padding=1)
        self.up2 = nn.Conv1d(in_channels=filters[2], out_channels=filters[1], kernel_size=3,padding=1)
        self.up1 = nn.Conv1d(in_channels=filters[1], out_channels=filters[0], kernel_size=3,padding=1)
        self.end = nn.Conv1d(in_channels=filters[0], out_channels=num_output_channels, kernel_size=3,padding=1)

        self.pool = nn.MaxPool1d(kernel_size=2,stride=2)
        self.upsample = nn.Upsample(scale_factor=2, mode='nearest')


    def forward(self, x):
        x = x.T.unsqueeze(0)
        x = self.start(x)
        x0 = x.clone() # 保存卷积层特征图，用于跳跃连接
        x = self.pool(x)
        x = self.down1(x)
        x1 = x.clone() # 保存卷积层特征图，用于跳跃连接
        x = self.pool(x)
        x = self.down2(x)
        x2 = x.clone() # 保存卷积层特征图，用于跳跃连接
        x = self.pool(x)
        x = self.down3(x)
        x3 = x.clone() # 保存卷积层特征图，用于跳跃连接
        x = self.pool(x)
        x = self.down4(x)
        x4 = x.clone() # 保存卷积层特征图，用于跳跃连接
        # 开始上采样以及跳连
        x = self.up4(x+x4)
        x = self.upsample(x)
        x = self.up3(x+x3)
        x = self.upsample(x)
        x = self.up2(x+x2)
        x = self.upsample(x)
        x = self.up1(x+x1)
        x = self.upsample(x)
        # 输出层
        x = self.end(x)
        return x.squeeze(0).permute(1, 0)






